from tensorflow.keras.utils import Sequence
import numpy as np
from skimage.io import imread
import skimage
import os

from layer_utils import anchor_boxes, get_gt_data
import layer_utils


class DataGenerator(Sequence):
    def __init__(self, args, batch_size, dictionary, n_classes, feature_shapes,
                 n_anchors, shuffle) -> None:
        super().__init__()
        self.args = args
        self.dictionary = dictionary
        self.n_classes = n_classes
        self.feature_shapes = feature_shapes
        self.n_anchors = n_anchors
        self.shuffle = shuffle
        self.batch_size = batch_size

        self.keys = np.array(list(self.dictionary.keys()))
        self.input_shape = (args.height, args.width, args.channels)

        self.on_epoch_end()
        self.get_n_boxes()

    def on_epoch_end(self):
        if self.shuffle == True:
            np.random.shuffle(self.keys)

    def __len__(self):
        blen = np.floor(len(self.dictionary) / self.batch_size)
        return int(blen)

    def __getitem__(self, index):
        start_index = index * self.batch_size
        end_index = (index + 1) * self.batch_size
        keys = self.keys[start_index:end_index]
        x, y = self.__data_generation(keys)
        return x, y

    def get_n_boxes(self):
        self.n_boxes = 0
        for shape in self.feature_shapes:
            self.n_boxes += np.prod(shape) // self.n_anchors
        return self.n_boxes

    def __data_generation(self, keys):
        x = np.zeros((self.batch_size, *self.input_shape))
        dim = (self.batch_size, self.n_boxes, self.n_classes)
        gt_class = np.zeros(dim)
        dim = (self.batch_size, self.n_boxes, 4)
        gt_offset = np.zeros(dim)
        gt_mask = np.zeros(dim)

        for i, key in enumerate(keys):
            video_id, pic_id = key.split('-')
            image_path = os.path.join(self.args.data_path, 'train_images',
                                      'video_' + video_id, pic_id + '.jpg')
            if self.args.no_img:
                image = np.zeros(self.input_shape)
            else:
                image = skimage.img_as_float(imread(image_path))
            x[i] = image

            labels = self.dictionary[key]
            labels = np.array(labels)

            boxes = labels[:, 0:-1]
            for index, feature_shape in enumerate(self.feature_shapes):
                anchors = anchor_boxes(
                    feature_shape, image.shape, index, self.args.layers)
                anchors = np.reshape(anchors, [-1, 4])
                iou = layer_utils.iou(anchors, boxes)

                gt = get_gt_data(iou, self.n_classes, anchors,
                                 labels, self.args.threshold)

                gt_cls, gt_off, gt_msk = gt

                if index == 0:
                    cls = np.array(gt_cls)
                    off = np.array(gt_off)
                    msk = np.array(gt_msk)
                else:
                    cls = np.append(cls, gt_cls, axis=0)
                    off = np.append(off, gt_off, axis=0)
                    msk = np.append(msk, gt_msk, axis=0)
            gt_class[i] = cls
            gt_offset[i] = off
            gt_mask[i] = msk

        y = [gt_class, np.concatenate((gt_offset, gt_mask), axis=-1)]

        return x, y
